import time
from drill_script.generate import generate_drill_script, generate_drill_script_stream
from drill_script.query_rewrite import query_rewrite, query_rewrite_retrieval_only
from drill_script.recall import recall
from drill_script.rerank import rerank


class DrillScriptPipeline:
    """演练脚本生成pipeline"""
    
    def __init__(self, equipment: str, scene_desc: str, plan: dict, script_node: list, retrieval_only=False):
        """
        初始化pipeline
        
        Args:
            equipment: 设备
            scene_desc: 场景描述
            plan: 应急预案
            script_node: 原始脚本节点结构
            retrieval_only: 是否使用纯检索模式
        """
        self.equipment = equipment
        self.scene_desc = scene_desc
        self.plan = plan
        self.script_node = script_node
        self.retrieval_only = retrieval_only  # 新增参数：是否使用纯检索模式
    
    async def process(self):
        """
        处理演练脚本生成
        
        Returns:
            生成的演练脚本或检索到的相关文档
        """
        # 记录pipeline开始时间
        pipeline_start_time = time.time()
        
        print(f"DrillScriptPipeline开始处理")
        print(f"设备: {self.equipment}")
        print(f"场景描述: {self.scene_desc}")
        print(f"预案字段数: {len(self.plan)}")
        print(f"脚本节点数: {len(self.script_node)}")
        print(f"处理模式: {'纯检索模式' if self.retrieval_only else 'RAG模式'}")
        
        try:
            # 记录参数验证开始时间
            validation_start_time = time.time()
            
            # 参数验证
            if not self.equipment or not self.scene_desc:
                raise Exception("设备或场景描述不能为空")
            
            if not self.plan or not isinstance(self.plan, dict):
                raise Exception("应急预案不能为空且必须是字典格式")
            
            if not self.script_node or not isinstance(self.script_node, list):
                raise Exception("脚本节点不能为空且必须是列表格式")
            
            validation_end_time = time.time()
            validation_duration = validation_end_time - validation_start_time
            print(f"参数验证耗时: {validation_duration:.3f}秒")
            
            # 查询重写
            print(f"开始查询重写...")
            rewrite_start_time = time.time()
            if self.retrieval_only:
                # 使用纯检索模式
                rewrite_result = await query_rewrite_retrieval_only(self.equipment, self.scene_desc, self.plan)
            else:
                # 使用RAG模式（大模型）
                rewrite_result = await query_rewrite(self.equipment, self.scene_desc, self.plan)
            rewrite_end_time = time.time()
            rewrite_duration = rewrite_end_time - rewrite_start_time
            print(f"查询重写完成: {rewrite_result}")
            print(f"查询重写耗时: {rewrite_duration:.3f}秒")
            
            # 召回相关文档
            print(f"开始文档召回...")
            recall_start_time = time.time()
            recall_result = recall(rewrite_result)
            recall_end_time = time.time()
            recall_duration = recall_end_time - recall_start_time
            print(f"文档召回完成，召回数量: {len(recall_result) if recall_result else 0}")
            print(f"文档召回耗时: {recall_duration:.3f}秒")
            
            # 重排序
            print(f"开始文档重排序...")
            rerank_start_time = time.time()
            rerank_result = rerank(recall_result)
            rerank_end_time = time.time()
            rerank_duration = rerank_end_time - rerank_start_time
            print(f"文档重排序完成，重排序数量: {len(rerank_result) if rerank_result else 0}")
            print(f"文档重排序耗时: {rerank_duration:.3f}秒")
            
            if self.retrieval_only:
                # 纯检索模式：直接返回检索到的文档
                print(f"纯检索模式：直接返回检索到的相关文档")
                
                # 构建返回结果 - 保持与大模型模式一致的结构
                result = {
                    "scriptNode": self.script_node,  # 直接返回原始script_node，与大模型模式保持一致
                    "related_documents": []
                }
                
                # 将检索到的文档转换为可序列化的格式
                for i, article in enumerate(rerank_result):
                    doc_info = {
                        "rank": i + 1,
                        "title": article.name,
                        "content": article.content,
                        "file_id": article.file_id,
                        "score": article.score,
                        "permission_type": article.permission_type,
                        "permission_set": article.permission_set
                    }
                    result["related_documents"].append(doc_info)
                
                # 记录pipeline结束时间
                pipeline_end_time = time.time()
                pipeline_duration = pipeline_end_time - pipeline_start_time
                
                print(f"演练脚本检索完成，返回 {len(result['related_documents'])} 个相关文档")
                print(f"Pipeline总体耗时: {pipeline_duration:.3f}秒")
                print(f"Pipeline耗时分析:")
                print(f"   - 参数验证: {validation_duration:.3f}秒 ({(validation_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 查询重写: {rewrite_duration:.3f}秒 ({(rewrite_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 文档召回: {recall_duration:.3f}秒 ({(recall_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 文档重排序: {rerank_duration:.3f}秒 ({(rerank_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 结果构建: {pipeline_duration - validation_duration - rewrite_duration - recall_duration - rerank_duration:.3f}秒")
                
                return result
            else:
                # RAG模式：使用大模型生成演练脚本
                print(f"RAG模式：开始生成演练脚本...")
                generation_start_time = time.time()
                result = await generate_drill_script(
                    self.equipment,
                    self.scene_desc,
                    self.plan,
                    self.script_node,
                    rerank_result  # 传递相关文档
                )
                generation_end_time = time.time()
                generation_duration = generation_end_time - generation_start_time
                print(f"演练脚本生成完成")
                print(f"脚本生成耗时: {generation_duration:.3f}秒")
                
                # 记录pipeline结束时间
                pipeline_end_time = time.time()
                pipeline_duration = pipeline_end_time - pipeline_start_time
                
                print(f"Pipeline总体耗时: {pipeline_duration:.3f}秒")
                print(f"Pipeline耗时分析:")
                print(f"   - 参数验证: {validation_duration:.3f}秒 ({(validation_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 查询重写: {rewrite_duration:.3f}秒 ({(rewrite_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 文档召回: {recall_duration:.3f}秒 ({(recall_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 文档重排序: {rerank_duration:.3f}秒 ({(rerank_duration/pipeline_duration)*100:.1f}%)")
                print(f"   - 脚本生成: {generation_duration:.3f}秒 ({(generation_duration/pipeline_duration)*100:.1f}%)")
                
                return result
            
        except Exception as e:
            print(f"❌ DrillScriptPipeline处理失败: {e}")
            import traceback
            traceback.print_exc()
            raise e
    
    async def process_stream(self):
        """
        流式处理演练脚本生成
        
        Yields:
            生成过程中的内容片段
        """
        print(f"DrillScriptPipeline开始流式处理")
        print(f"设备: {self.equipment}")
        print(f"场景描述: {self.scene_desc}")
        print(f"预案字段数: {len(self.plan)}")
        print(f"脚本节点数: {len(self.script_node)}")
        print(f"处理模式: {'纯检索模式' if self.retrieval_only else 'RAG模式'}")
        
        try:
            # 参数验证
            if not self.equipment or not self.scene_desc:
                yield f"错误: 设备或场景描述不能为空"
                return
            
            if not self.plan or not isinstance(self.plan, dict):
                yield f"错误: 应急预案不能为空且必须是字典格式"
                return
            
            if not self.script_node or not isinstance(self.script_node, list):
                yield f"错误: 脚本节点不能为空且必须是列表格式"
                return
            
            # 查询重写
            yield f"开始查询重写...\n"
            if self.retrieval_only:
                rewrite_result = await query_rewrite_retrieval_only(self.equipment, self.scene_desc, self.plan)
            else:
                rewrite_result = await query_rewrite(self.equipment, self.scene_desc, self.plan)
            yield f"查询重写完成，关键词: {rewrite_result}\n"
            
            # 召回相关文档
            yield f"开始文档召回...\n"
            recall_result = recall(rewrite_result)
            yield f"文档召回完成，召回数量: {len(recall_result) if recall_result else 0}\n"
            
            # 重排序
            yield f"开始文档重排序...\n"
            rerank_result = rerank(recall_result)
            yield f"文档重排序完成，重排序数量: {len(rerank_result) if rerank_result else 0}\n"
            
            if self.retrieval_only:
                # 纯检索模式：直接返回检索结果
                yield f"纯检索模式：直接返回检索到的相关文档\n"
                
                # 构建返回结果
                result = {
                    "status": "success",
                    "message": "检索完成，返回相关演练脚本文档",
                    "mode": "retrieval_only",
                    "data": {
                        "equipment": self.equipment,
                        "scene_desc": self.scene_desc,
                        "plan": self.plan,
                        "script_node": self.script_node,
                        "related_documents": []
                    }
                }
                
                # 将检索到的文档转换为可序列化的格式
                for i, article in enumerate(rerank_result):
                    doc_info = {
                        "rank": i + 1,
                        "title": article.name,
                        "content": article.content,
                        "file_id": article.file_id,
                        "score": article.score,
                        "permission_type": article.permission_type,
                        "permission_set": article.permission_set
                    }
                    result["data"]["related_documents"].append(doc_info)
                
                yield f"演练脚本检索完成，返回 {len(result['data']['related_documents']) if result['data']['related_documents'] else 0} 个相关文档\n"
                yield f"检索结果: {result}\n"
                
                print(f"演练脚本流式检索完成")
            else:
                # RAG模式：流式生成演练脚本
                yield f"RAG模式：开始流式生成演练脚本...\n"
                async for chunk in generate_drill_script_stream(
                    self.equipment,
                    self.scene_desc,
                    self.plan,
                    self.script_node
                ):
                    yield chunk
                
                print(f"演练脚本流式生成完成")
            
        except Exception as e:
            print(f"DrillScriptPipeline流式处理失败: {e}")
            import traceback
            traceback.print_exc()
            yield f"处理失败: {e}" 